#include "Webcam.h"

#include <atlbase.h>
#include <atlstr.h>
#include <windows.h>
#include <dshow.h>
#include <streams.h> // Baseclasses library

#include "Log.h"

// {38487B2E-5431-4187-B354-49B5955569E4}
static const GUID CLSID_JPCUSTOMRENDERER = 
{ 0x38487b2e, 0x5431, 0x4187, { 0xb3, 0x54, 0x49, 0xb5, 0x95, 0x55, 0x69, 0xe4 } };

class CustomRenderer : public CBaseRenderer 
{
private:
	VIDEOINFOHEADER m_currentVideoInfoHeader;

	HRESULT m_baseRendererHr;

	WebcamFrame m_currentFrame;

	Webcam* m_webcam;

public:

	CustomRenderer(IUnknown* owner = 0)
		: CBaseRenderer(CLSID_JPCUSTOMRENDERER, L"CustomRenderer", owner, &m_baseRendererHr)
	{
		m_currentFrame.data = 0;
	}

	virtual ~CustomRenderer()
	{
		if(m_currentFrame.data)
		{
			delete [] m_currentFrame.data;
		}

		// FIXME: This is just wrong..
		if(m_pInputPin)
		{
			//m_pInputPin->Disconnect();
			//m_pInputPin->Release();
			//m_pInputPin->Release();

			m_pInputPin = 0;
		}
	}

	virtual HRESULT CheckMediaType(const CMediaType* mediaType)
	{
		AM_MEDIA_TYPE* type = (AM_MEDIA_TYPE*)mediaType;

		if(type->majortype == MEDIATYPE_Video)
		{
			if(type->subtype == MEDIASUBTYPE_RGB32)
			{
				if(type->formattype == FORMAT_VideoInfo)
				{
					AM_MEDIA_TYPE* amMediaType = (AM_MEDIA_TYPE*)mediaType;
					VIDEOINFOHEADER* header = (VIDEOINFOHEADER*)amMediaType->pbFormat;

					// Use a a format that's as close or better to the preferred size
					if(header->bmiHeader.biWidth >= m_webcam->preferredWidth() && header->bmiHeader.biHeight >= m_webcam->preferredHeight())
					{
						m_currentVideoInfoHeader = *header;
						return S_OK;
					}
				}
			}
	
		}

		return S_FALSE;
	}

	virtual HRESULT DoRenderSample(IMediaSample* pMediaSample)
	{
		AM_MEDIA_TYPE* mt;
		pMediaSample->GetMediaType(&mt);

		// TODO: Fix this. mt should be valid during a video mode switch
		//		 and invalid anytime else. Figure out how to handle this
		//		 gracefully.

		VIDEOINFOHEADER* header = &m_currentVideoInfoHeader;//(VIDEOINFOHEADER*)mt->pbFormat;

		int width = header->bmiHeader.biWidth;
		int height = header->bmiHeader.biHeight;
		int bpp = header->bmiHeader.biBitCount;

		unsigned char* data;
		pMediaSample->GetPointer(&data);

		m_currentFrame.mediaType.width = width;
		m_currentFrame.mediaType.height = height;
		m_currentFrame.mediaType.format = WebcamMediaType::Rgb32;// = bpp; // TODO: Change this if need be?

		if(!m_currentFrame.data)
		{
			m_currentFrame.data = new unsigned char[width * height * (bpp >> 3)];
		}

		memcpy(m_currentFrame.data, data, width * height * (bpp >> 3));

		m_webcam->setCurrentFrame(&m_currentFrame);

		return S_OK;
	}

	virtual HRESULT CompleteConnect(IPin *pReceivePin)
	{
		HRESULT res = CBaseRenderer::CompleteConnect(pReceivePin);
		return res;
	}

	void setWebcam(Webcam* webcam)
	{
		m_webcam = webcam;
	}
};

class Webcam::WebcamImpl
{
public:

	WebcamImpl(Webcam* owner)
		: m_owner(owner)
	{ }

	~WebcamImpl()
	{
		delete m_renderer;
		
		
		if(m_mediaControl)
		{
			m_mediaControl->Release();
			m_mediaControl = 0;
		}

		if(m_captureFilter)
		{
			m_captureFilter->Release();
			m_captureFilter = 0;
		}

		if(m_captureGraph)
		{
			m_captureGraph->Release();
			m_captureGraph = 0;
		}

		if(m_filterGraph)
		{
			m_filterGraph->Release();
			m_filterGraph = 0;
		}
		
	}

	bool setupFilterGraph()
	{
		HRESULT result = CoCreateInstance(CLSID_FilterGraph, 0, CLSCTX_INPROC_SERVER, IID_IGraphBuilder, reinterpret_cast<void**>(&m_filterGraph));
		if(FAILED(result))
		{
			log::debug << "Webcam::WebcamImpl::setupFilterGraph(): [Error] - Failed to create filter graph instance.\n";
			return false;
		}

		int t = m_filterGraph->AddRef();

		return true;
	}

	bool setupCaptureGraph()
	{
		if(!m_filterGraph)
		{
			log::debug << "Webcam::WebcamImpl::setupCaptureGraph(): [Error] - Initialize the filter graph first.\n";
			return false;
		}

		HRESULT result = CoCreateInstance(CLSID_CaptureGraphBuilder2, 0, CLSCTX_INPROC_SERVER, IID_ICaptureGraphBuilder2, reinterpret_cast<void**>(&m_captureGraph));
		if(FAILED(result))
		{
			log::debug << "Webcam::WebcamImpl::setupCaptureGraph(): [Error] - Failed to create camera graph instance.\n";
			return false;
		}

		int t = m_captureGraph->AddRef();
		m_captureGraph->SetFiltergraph(m_filterGraph);

		return true;
	}

	bool setupInterfaces()
	{
		HRESULT result = m_filterGraph->QueryInterface(IID_IMediaControl, reinterpret_cast<void**>(&m_mediaControl));
		if(FAILED(result))
		{
			log::debug << "Webcam::WebcamImpl::setupInterfaces(): [Error] - Failed to query media control interface.\n";
			return false;
		}

		return true;
	}

	bool setCaptureFilter(IBaseFilter* filter)
	{
		HRESULT result = m_filterGraph->AddFilter(filter, L"CaptureFilter");
		if(FAILED(result))
		{
			log::debug << "Webcam::WebcamImpl::setCaptureFilter(): [Error] - Could not add capture filter.\n";
			return false;
		}

		m_captureFilter = filter;
		int t = m_captureFilter->AddRef();

		return true;
	}

	bool setDevice(const WebcamDevice& device)
	{
		CComPtr<IBaseFilter> filter = findDeviceFilterByName(device.name);
		if(!filter)
		{
			log::debug << "Webcam::WebcamImpl::setDevice(): [Error] - Invalid device filter.\n";
			return false;
		}

		return setCaptureFilter(filter);
	}

	bool connectToRenderer()
	{
		CComPtr<IFilterGraph2> filterGraph2;
		m_filterGraph->QueryInterface(IID_IFilterGraph2, reinterpret_cast<void**>(&filterGraph2));

		CComPtr<IEnumPins> enumPins;
		m_captureFilter->EnumPins(&enumPins);

		log::debug << "Listing pins... \n";

		IPin* pin;

		int pinCount = 0;

		while(enumPins->Next(1, &pin, 0) == S_OK)
		{
			// Get the first pin
			if(pinCount == 0)
				pin = pin;

			PIN_INFO info;
			pin->QueryPinInfo(&info);

			log::debug << "Found " << ((info.dir == PINDIR_INPUT) ? "Input" : "Output") <<  " Pin: " << info.achName << ".\n";

			pin->Release();

			++pinCount;
		}

		log::debug << "Found " << pinCount << " pin(s).\n";

		m_renderer = new CustomRenderer();
		m_renderer->setWebcam(m_owner);

		HRESULT res = 0;

		res = m_filterGraph->AddFilter(m_renderer, L"CustomRenderer");
		if(FAILED(res))
		{
			log::debug << "Webcam::WebcamImpl::connectToRenderer(): [Error] - Failed to add render filter.\n";
			return false;
		}
		
		res = m_captureGraph->RenderStream(&PIN_CATEGORY_CAPTURE, &MEDIATYPE_Video, m_captureFilter, 0, m_renderer);
		if(FAILED(res))
		{
			log::debug << "Webcam::WebcamImpl::connectToRenderer(): [Error] - Failed to connect stream pins.\n";
			return false;
		}
		
		return true;
	}

	bool run()
	{
		if(!connectToRenderer())
			return false;

		HRESULT res;
		//while((res = m_mediaControl->Run()) == S_FALSE);
		res = m_mediaControl->Run();

		//FILTER_STATE state;
		//m_renderer->GetState(0, &state);

		return true;
	}

	void stop()
	{
		m_mediaControl->Stop();
	}

	Webcam* owner()
	{
		return m_owner;
	}

private:
	IGraphBuilder* m_filterGraph;
	ICaptureGraphBuilder2* m_captureGraph;
	IBaseFilter* m_captureFilter;
	IMediaControl* m_mediaControl;

	CustomRenderer* m_renderer;

	Webcam* m_owner;

	CComPtr<IBaseFilter> findDeviceFilterByName(const std::string& deviceName)
	{
		HRESULT result;

		CComPtr<ICreateDevEnum> deviceEnum;
		result = CoCreateInstance(CLSID_SystemDeviceEnum, 0, CLSCTX_INPROC_SERVER, IID_ICreateDevEnum, reinterpret_cast<void**>(&deviceEnum));
		if(FAILED(result))
		{
			log::debug << "Webcam::WebcamImpl::findDeviceFilterByName(): [Error] - Failed to create device enumerator.\n";
			return false;
		}

		CComPtr<IEnumMoniker> classEnumMoniker;
		result = deviceEnum->CreateClassEnumerator(CLSID_VideoInputDeviceCategory, &classEnumMoniker, 0);
		if(FAILED(result))
		{
			log::debug << "Webcam::WebcamImpl::findDeviceFilterByName(): [Error] - Failed to create video input device enumerator.\n";
			return false;
		}

		if(!classEnumMoniker)
		{
			log::debug << "Webcam::WebcamImpl::findDeviceFilterByName(): [Error] - Your system does not have any video input devices.\n";
			return false;
		}

		CComPtr<IBaseFilter> filter;

		IMoniker* moniker;
		while(classEnumMoniker->Next(1, &moniker, 0) == S_OK)
		{
			CComPtr<IPropertyBag> properties = 0;
			result = moniker->BindToStorage(0, 0, IID_IPropertyBag, reinterpret_cast<void**>(&properties));
			if(FAILED(result))
				continue;

			VARIANT var;
			var.vt = VT_BSTR;
			properties->Read(L"FriendlyName", &var, 0);
			std::string name = UnicodeToASCII(var.bstrVal);

			if(name == deviceName)
			{
				result = moniker->BindToObject(0, 0, IID_IBaseFilter, reinterpret_cast<void**>(&filter));
			}

			moniker->Release();
		}

		if(!filter)
			log::debug << "Webcam::WebcamImpl::findDeviceFilterByName(): [Error] - Did not find device by name \'" << deviceName << "\'.\n";

		return filter;
	}
};

Webcam::Webcam()
: m_impl(new Webcam::WebcamImpl(this)), m_currentFrame(0)
{
	CoInitialize(0);

	setPreferredSize(640, 480);
}

Webcam::~Webcam()
{
	delete m_impl;
	CoUninitialize();
}

bool Webcam::setup(const WebcamDevice& device)
{
	if(!m_impl->setupFilterGraph())
		return false;

	if(!m_impl->setupCaptureGraph())
		return false;

	if(!setDevice(device))
		return false;

	if(!m_impl->setupInterfaces())
		return false;

	return true;
}

bool Webcam::start()
{
	return m_impl->run();
}

void Webcam::stop()
{
	m_impl->stop();
}

const WebcamFrame* Webcam::currentFrame() const
{
	return m_currentFrame;
}

void Webcam::setCurrentFrame(WebcamFrame* currentFrame)
{
	m_currentFrame = currentFrame;
}

bool Webcam::setDevice(const WebcamDevice& device)
{
	m_device = device;
	m_impl->setDevice(m_device);
	return true;
}

void Webcam::setPreferredSize(int width, int height)
{
	m_preferredWidth = width;
	m_preferredHeight = height;
}

int Webcam::preferredWidth() const
{
	return m_preferredWidth;
}

int Webcam::preferredHeight() const
{
	return m_preferredHeight;
}

bool Webcam::listDevices(std::vector<WebcamDevice>& outDevices)
{
	HRESULT result;

	CComPtr<ICreateDevEnum> deviceEnum;
	result = CoCreateInstance(CLSID_SystemDeviceEnum, 0, CLSCTX_INPROC_SERVER, IID_ICreateDevEnum, reinterpret_cast<void**>(&deviceEnum));
	if(FAILED(result))
	{
		log::debug << "Webcam::listDevices(): [Error] - Failed to create device enumerator.\n";
		return false;
	}

	CComPtr<IEnumMoniker> classEnumMoniker;
	result = deviceEnum->CreateClassEnumerator(CLSID_VideoInputDeviceCategory, &classEnumMoniker, 0);
	if(FAILED(result))
	{
		log::debug << "Webcam::listDevices(): [Error] - Failed to create video input device enumerator.\n";
		return false;
	}

	if(!classEnumMoniker)
	{
		log::debug << "Webcam::listDevices(): [Error] - Your system does not have any video input devices.\n";
		return false;
	}

	IMoniker* moniker;
	while(classEnumMoniker->Next(1, &moniker, 0) == S_OK)
	{
		CComPtr<IPropertyBag> properties = 0;
		result = moniker->BindToStorage(0, 0, IID_IPropertyBag, reinterpret_cast<void**>(&properties));
		if(FAILED(result))
			continue;

		WebcamDevice device;
		
		VARIANT var;
		var.vt = VT_BSTR;
		properties->Read(L"FriendlyName", &var, 0);
		device.name = UnicodeToASCII(var.bstrVal);

		properties->Read(L"Description", &var, 0);
		device.description = UnicodeToASCII(var.bstrVal);
		outDevices.push_back(device);

		moniker->Release();
	}

	return true;
}
